module efield

  use dimensions
  use fft_3d

  real,    private   :: kk3D(jz,jy,jx)      ! matrix of k-squared
  complex, private   :: phi0init(jz,jy,jx)  ! initial potential of background Efield 
  complex, private   :: phi03DFT(jz,jy,jx)  ! FT of potential of background Efield


  complex, private   :: E0Cinit(jz)         ! initial background E-field at the center
  complex, private   :: E0C(jz)             ! background E-field at the center
  complex, private   :: PhiInt
  complex, private   :: PhiTot



contains

!--------------------------------------------------------------------------------------
subroutine efield_init()

  use fft_wrap, only   : fft_forward, fft_backward
  use parameters, only : phi0, kz, iE0Modulation, SpeckleWidth, SpeckleDepth, iunit
  use mpi_wrap

  implicit none

  integer              :: ix, iy, iz

  complex              :: xMod(jx)
  complex              :: yMod(jy)

  write(msg,*) char(10), ' Initialing efield module'
  call mpi_wrap_msg

  !-- setting modulation vectors --

  if (iE0Modulation) then

      xMod(:) = exp( - 20 * xvec(:)**2 / SpeckleWidth**2 )
      yMod(:) = exp( - 20 * yvec(:)**2 / SpeckleDepth**2 )

      call fft_forward(jx,xMod,xMod)
      xMod=xMod/float(jx)
      call deAlias(jx,xMod)
      call fft_backward(jx,xMod,xMod)  

      call fft_forward(jy,yMod,yMod)
      yMod=yMod/float(jy)
      call deAlias(jy,yMod)
      call fft_backward(jy,yMod,yMod)  

  else

     xMod(:) = 1.
     yMod(:) = 1.

  end if


  !-- initialization of background field 

  E0Cinit(:) =  -iunit * kz * phi0 *exp ( iunit * kz * zvec(:) )

  do ix=1,jx
     do iy=1,jy
        phi0init(:,iy,ix) = xMod(ix)*yMod(iy)*phi0 *exp ( iunit * kz * zvec(:)  ) 
     enddo
  enddo
  

  !-- take Fourier transform of initial potential, store result

  E0C      = E0Cinit
  phiTot   = phi0*kz*kz

  phi03DFT = phi0init

  call fft_forward_x(phi03DFT)
  call fft_forward_y(phi03DFT)
  call fft_forward_z(phi03DFT)


  !-- generate matrix of k-squared

  do ix=1,jx
     do iy=1,jy
        do iz=1,jz
           kk3D(iz,iy,ix) = kXvec(ix)**2 + kZvec(iz)**2  + kYvec(iy)**2
        end do
     end do
  end do

end subroutine efield_init


!-------------------------------------------------------------------------------------
subroutine efield_compute_E(dens, filter, &                       ! inputs 
                            Ez3D, Ex3D, Ey3D, OmegaBounce)        ! outputs
!-------------------------------------------------------------------------------------

  use fft_wrap,   only : fft_forward, fft_backward
  use parameters, only : iunit, nz

  implicit none


  !-- external variables --

  real,    intent(in)  ::  dens(jz,jy,jx) 
  real,    intent(in)  ::  filter(jz,jy,jx)

  real,    intent(out) ::  Ez3D(jz,jy,jx)
  real,    intent(out) ::  Ex3D(jz,jy,jx)
  real,    intent(out) ::  Ey3D(jz,jy,jx)
  real,    intent(out) ::  OmegaBounce


  !-- internal variables --

  complex  phi3DFT(jz,jy,jx)
  complex  dummy3D(jz,jy,jx)
  complex  dummy1Dz(jz)

  real     PhiCenter
  integer  ix, iy, iz

      
  !-- find FT of potential --

  dummy3D = dens
 
  call fft_forward_z(dummy3D)
  call fft_forward_x(dummy3D)
  call fft_forward_y(dummy3D)

  dummy3D(:,:,:) = dummy3D(:,:,:)*filter(:,:,:)

  phi3DFT = dummy3D / kk3D             ! inversing Laplacian in k-space


! phi3DFT(1,:,:) = 0.0        ! comment out for comparibility with BGK code
  phi3DFT(1,1,1) = 0.0        ! uncomment for comparibility with BGK code
 

  phi3DFT = phi3DFT + phi03DFT         ! total potential

  !phi3DFT(:,1,1) = 0                  ! spatially uniform in x not allowed (03/21/2011)
        
 

  !-- differentiate potential to find z-component of electric field -- 

  do iz=1,jz
     dummy3D(iz,:,:) = -iunit * kZvec(iz) * phi3DFT(iz,:,:)
  end do

  call fft_backward_z(dummy3D)
  call fft_backward_x(dummy3D)
  call fft_backward_y(dummy3D)

  Ez3D = dummy3D



  !-- differentiate potential to find x-component of electric field -- 

  if (jx .eq. 1) then

     Ex3D = 0.

  else

     do ix=1,jx
        dummy3D(:,:,ix) = -iunit * kXvec(ix) * phi3DFT(:,:,ix)
     end do

     call fft_backward_z(dummy3D)
     call fft_backward_x(dummy3D)
     call fft_backward_y(dummy3D)

     Ex3D = dummy3D

  end if


  !-- differentiate potential to find y-component of electric field -- 

  if (jy .eq. 1) then

     Ey3D = 0.

  else

     do iy=1,jy
        dummy3D(:,iy,:) = -iunit * kYvec(iy) * phi3DFT(:,iy,:)
     end do

     call fft_backward_z(dummy3D)
     call fft_backward_x(dummy3D)
     call fft_backward_y(dummy3D)

     Ey3D = dummy3D

  end if


  !-- find phi(center) for diagnostics later --

  dummy1Dz = dens(:, jy/2+1, jx/2+1)

  call fft_forward(jz, dummy1Dz, dummy1Dz)
  dummy1Dz = dummy1Dz/float(jz)
  call dealias(jz, dummy1Dz)

  phiInt = 2*dummy1Dz(1+nz)

  dummy1Dz(:) = Ez3D(:, jy/2+1, jx/2+1)

  call fft_forward(jz, dummy1Dz, dummy1Dz)
  dummy1Dz = dummy1Dz/float(jz)
  call dealias(jz, dummy1Dz)


  ! amplitude of TOTAL potential at speckle center of sin(kz*z + phase)
  PhiCenter = 2 * abs(dummy1Dz(nz+1) / kZvec(nz+1))

  OmegaBounce = kZvec(nz+1) * sqrt(PhiCenter)
     
  return

end subroutine efield_compute_E
    

!-------------------------------------------------------------------------------------
subroutine efield_evolve_E0(t)                                   ! input
!-------------------------------------------------------------------------------------

  use parameters, only : kz, phi0, vPhase, iWaveFrame, iunit

  implicit none

  real, intent(in)  :: t    
  complex           :: ExpKVt

  if ( iWaveFrame==1 ) return


  ExpKVt = exp( - iunit*t*vPhase*kz  )
     
  E0C      = ExpKVt * E0Cinit         
  phi03DFT = ExpKVt * phi0init
  phiTot   = ExpKVt * phi0*kz*kz
     
  !-- take Fourier transform of initial potential, store result

  call fft_forward_x(phi03DFT)
  call fft_forward_y(phi03DFT)
  call fft_forward_z(phi03DFT)


end subroutine efield_evolve_E0


!-------------------------------------------------------------------------------------
subroutine efield_get_epsilon(epsilon)                ! output 
!-------------------------------------------------------------------------------------

  implicit none

  complex, intent(out)    :: epsilon

  epsilon = phiTot / (phiInt + phiTot)
 
end subroutine efield_get_epsilon

!-------------------------------------------------------------------------------------
subroutine efield_get_E0z_center(E0zCenter)            ! output
!-------------------------------------------------------------------------------------

  implicit none

  real, intent(out)       :: E0zCenter(jz)

  E0zCenter = real(E0C)

end subroutine efield_get_E0z_center

!-------------------------------------------------------------------------------------

  
end module
